import numpy as np
import pylab as plt
from mpl_toolkits.mplot3d import Axes3D 

"""
Copy from https://github.com/HERA-Team/uvtools/blob/main/uvtools/plot.py
"""

def data_mode(data, mode='abs'):
    """
    Apply filter to data according to a chosen plotting mode.

    Parameters
    ----------
    data : array_like
        Array of data to be plotted (normally complex floats).

    mode : str, optional
        Which transformation to apply to the data. Options are:
          - 'phs':  Phase angle.
          - 'abs':  Absolute value.
          - 'real': Real value.
          - 'imag': Imaginary value.
          - 'log':  Log (base-10) of absolute value.
        Default: 'abs'.

    Returns
    -------
    data : array_like
        Data transformed according to the value of `mode`.
    """
    if mode.startswith('phs'):
        data = np.angle(data)
    elif mode.startswith('abs'):
        data = np.absolute(data)
    elif mode.startswith('real'):
        data = data.real
    elif mode.startswith('imag'):
        data = data.imag
    elif mode.startswith('log'):
        data = np.absolute(data)
        data = np.log10(data)
    else:
        raise ValueError('Unrecognized plot mode.')
    return data

def waterfall(data, mode='log', vmin=None, vmax=None, drng=None, mx=None, **kwargs):
    """
    Generate a 2D waterfall plot
    Args:
        data(array-like): 2D complex array of data.

        mode : str, optional
            Which transformation to apply to the data. Options are:
            - 'phs':  Phase angle.
            - 'abs':  Absolute value.
            - 'real': Real value.
            - 'imag': Imaginary value.
            - 'log':  Log (base-10) of absolute value.
            Default: 'abs'.
        
        vmin, vmax : float, optional
            Minimum and maximum values of the color scale. If not set (and `mx` and
            `drng` are not set), the min. and max. values of the data will be used.

            Note that that the min. and max. values are the ones _after_ the data
            have been transformed according to `mode`. So, if `mode='log'`, these
            values are the min. and max. of log_10(data).

        mx : float, optional
            The max. value of the color scale in the plot (equivalent to vmax).
            Cannot be specified at the same time as `vmin` and `vmax`.

        drng : float, optional
            The difference between the min. and max. values of the color scale in
            the plot, `drng = mx - min`, where these are the min/max values after
            applying the transformation specified by `mode`.

            Cannot be specified at the same time as `vmin` and `vmax`.


    Returns
        plot : matplotlib.imshow
            Waterfall plot.
    """

    validity_msg = "Must specify either 'vim' and 'vmax' *or* 'mx' and 'drng'"
    if mx is not None or drng is not None:
        assert vmin is None and vmax is None,validity_msg
    if vmin is not None or vmax is not None:
        assert mx is None and drng is None,validity_msg
        mx = vmax
        drng = vmax - vmin

    # Fill masked array and recenter if requested
    if np.ma.isMaskedArray(data):
        data = data.filled(0)
    
    # if recenter:
    #     data = aipy.img.recenter(data, np.array(data.shape)/2)
    
    # Apply requested transform to data
    data = data_mode(data, mode=mode)

    # Get min/max values for color scale
    if mx is None:
        mx = data.max()
    if drng is None:
        drng = mx - data.min()
    mn = mx - drng

    # Choose aspect ratio
    if 'aspect' not in kwargs.keys():
        kwargs['aspect'] = 'auto'

    return plt.imshow(data, vmax=mx, vmin=mn, interpolation='nearest', **kwargs)


def waterfall3D(freq,time,data,size=(30,15)):
    """
    Generate a 2D waterfall plot
    Args:
        data(array-like): 2D complex array of data.
        size:the size of figure

    Returns
        plot : matplotlib.imshow
            Waterfall plot-3D.
    """
    x,y = np.meshgrid(freq,time)
    data = data_mode(data, mode="abs")
    fig = plt.figure(figsize=size)
    ax = Axes3D(fig)
    ax.set_xlabel('freq')
    ax.set_ylabel('time')
    ax.set_zlabel('frequency')
    return ax.plot_surface(x,y,data,cmap='rainbow')
